package com.lq;

import org.wltea.analyzer.cfg.Configuration;
import org.wltea.analyzer.core.IKSegmenter;
import org.wltea.analyzer.core.Lexeme;

import java.io.IOException;
import java.io.Reader;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.stream.Collectors;

public class IKSegmenterWrapper {

    private IKSegmenter smart;

    private final IKSegmenter impl;

    private Iterator<Lexeme> it;

    public IKSegmenterWrapper(Configuration configuration)
    {
        if (!configuration.isUseSmart()) {
            Configuration smartConfig = new Configuration(configuration.getEnvironment(), configuration.getSettings());
            smartConfig.setUseSmart(true);
            smart = new IKSegmenter(smartConfig);
        }
        impl = new IKSegmenter(configuration);
    }

    private void setData()
    {
        LinkedHashSet<Lexeme> all = new LinkedHashSet<>();
        try {
            Lexeme lexeme;
            while ((lexeme = impl.next()) != null) {
                all.add(lexeme);
            }
            if (smart != null) {
                while ((lexeme = smart.next()) != null) {
                    all.add(lexeme);
                }
            }
        } catch (IOException e) {
            throw new RuntimeException();
        }
        all = all.stream()
                .sorted(Lexeme::compareTo)
                .collect(Collectors.toCollection(LinkedHashSet::new));
        it = all.iterator();
    }

    public Lexeme next()
    {
        if (it == null) {
            return null;
        }
        if (it.hasNext()) {
            return it.next();
        }
        return null;
    }

    public void reset(Reader input) throws IOException
    {
        if (smart != null) {
            Reader[] inputs = CloneReader.clone(input);
            input = inputs[0];
            smart.reset(inputs[1]);
        }
        impl.reset(input);
        setData();
    }

}
